{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "92c9da5866bdcf7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T02:14:55.307750Z",
     "start_time": "2025-05-23T02:14:54.383318Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'年', '浮', '学', '翔', '文', '碧', '苍', '浪', '飞', '头', '意', '怅', '沉', '侯', '书', '户', '峥', '主', '天', '鹰', '记', '少', '\\n', '舟', '正', '透', '红', '鱼', '当', '侣', '风', '湘', '类', '立', '染', '往', '击', '问', '茂', '层', '底', '水', '寒', '点', '子', '竞', '地', '谁', '。', '曾', '嵘', '空', '万', '？', '由', '恰', '遏', '激', '月', '挥', '寥', '忆', '生', '否', '；', '粪', '独', '流', '来', '洲', '廓', '岁', '争', '游', '秋', '橘', '长', '土', '浅', '江', '尽', '稠', '山', '去', '，', '林', '气', '指', '到', '字', '中', '北', '自', '遒', '扬', '漫', '舸', '华', '方', '斥', '同', '百', '遍', '携', '看', '大', '霜', '昔', '茫'}\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "# 示例文本数据，一首诗\n",
    "text = \"\"\"\n",
    "独立寒秋，湘江北去，橘子洲头。\n",
    "看万山红遍，层林尽染；漫江碧透，百舸争流。\n",
    "鹰击长空，鱼翔浅底，万类霜天竞自由。\n",
    "怅寥廓，问苍茫大地，谁主沉浮？\n",
    "携来百侣曾游，忆往昔峥嵘岁月稠。\n",
    "恰同学少年，风华正茂；书生意气，挥斥方遒。\n",
    "指点江山，激扬文字，粪土当年万户侯。\n",
    "曾记否，到中流击水，浪遏飞舟？\n",
    "\"\"\"\n",
    "\n",
    "# 创建词汇表\n",
    "words = set(text)\n",
    "vocab_size = len(words)\n",
    "word_to_idx = {word: i for i, word in enumerate(words)}\n",
    "idx_to_word = {i: word for i, word in enumerate(words)}\n",
    "\n",
    "print(words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "91280acf83012c57",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T02:15:47.331640Z",
     "start_time": "2025-05-23T02:15:47.321464Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[22, 66, 33, 42, 74, 84, 31, 79, 91, 83, 84, 75, 44, 69, 9, 48, 22, 104, 52, 82, 26, 102, 84, 39, 85, 80, 34, 64, 95, 79, 5, 25, 84, 101, 96, 72, 67, 48, 22, 19, 36, 76, 51, 84, 27, 3, 78, 40, 84, 52, 32, 106, 18, 45, 92, 54, 48, 22, 11, 60, 70, 84, 37, 6, 108, 105, 46, 84, 47, 17, 12, 1, 53, 22, 103, 68, 101, 29, 49, 73, 84, 61, 35, 107, 16, 50, 71, 58, 81, 48, 22, 55, 100, 2, 21, 0, 84, 30, 97, 24, 38, 64, 14, 62, 10, 86, 84, 59, 99, 98, 93, 48, 22, 87, 43, 79, 82, 84, 57, 94, 4, 89, 84, 65, 77, 28, 0, 52, 15, 13, 48, 22, 49, 20, 63, 84, 88, 90, 67, 36, 41, 84, 7, 56, 8, 23, 53, 22]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# 超参数设置\n",
    "SEQ_LENGTH = 10  # 输入序列长度\n",
    "BATCH_SIZE = 1\n",
    "HIDDEN_SIZE = 128\n",
    "NUM_LAYERS= 1   # 隐藏层有2层，数据量少，层数多了反倒不好\n",
    "LEARNING_RATE = 0.005\n",
    "NUM_EPOCHS = 200\n",
    "\n",
    "\n",
    "# 创建训练数据\n",
    "class TextDataset(Dataset):\n",
    "    def __init__(self, text, seq_length):\n",
    "        self.text = text\n",
    "        self.seq_length = seq_length\n",
    "\n",
    "        # 转换为索引序列\n",
    "        self.data = [word_to_idx[ch] for ch in text]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data) - self.seq_length\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # 文本里的某个序列\n",
    "        input_seq = self.data[idx:idx + self.seq_length]\n",
    "\n",
    "        # 目标序列\n",
    "        target_seq = self.data[idx + 1:idx + self.seq_length + 1]\n",
    "\n",
    "        # 相当于，假如语料为abcdefg, input_seq=abc, target_seq=bcd\n",
    "\n",
    "        return torch.LongTensor(input_seq), torch.LongTensor(target_seq)\n",
    "\n",
    "\n",
    "dataset = TextDataset(text, SEQ_LENGTH)\n",
    "dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "\n",
    "print(dataset.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d9e9fd013d3d01ba",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T02:15:52.653582Z",
     "start_time": "2025-05-23T02:15:52.639686Z"
    }
   },
   "outputs": [],
   "source": [
    "class CharGRU(nn.Module):\n",
    "    def __init__(self, vocab_size, hidden_size):\n",
    "        super().__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        # 嵌入层（将词转为词向量）\n",
    "        # 词向量要输入给隐藏层，所以词向量的大小就是隐藏层的大小\n",
    "        self.embedding = nn.Embedding(vocab_size, hidden_size)\n",
    "\n",
    "        # RNN参数\n",
    "        # input_size表示输入的x的大小，一个词就是词向量，而这里词向量的大小就是hidden_size\n",
    "        self.rnn = nn.GRU(input_size=hidden_size, hidden_size=hidden_size, num_layers=NUM_LAYERS, batch_first=True)\n",
    "\n",
    "        # 输出层\n",
    "        self.out_linear = nn.Linear(hidden_size, vocab_size)\n",
    "\n",
    "    def forward(self, x, hidden=None):\n",
    "\n",
    "        # 初始化隐藏层状态\n",
    "        # 每送进来一个Batch就初始化一个hidden\n",
    "        if hidden is None:\n",
    "            hidden = torch.zeros(NUM_LAYERS, BATCH_SIZE, self.hidden_size)\n",
    "\n",
    "        # 嵌入层转换，x表示一个词序列，比如序列长度为10，那么embedded就是的形状就是(BATCH_SIZE, 10, hidden_size)，hidden_size就是词向量的大小（这里的词是一个汉字）\n",
    "        # 两维变成了三维\n",
    "        embedded = self.embedding(x)\n",
    "\n",
    "        # 存储所有时间步的输出\n",
    "        # 输入的形状决定了有多少个时间步，每个时间步的输出大小为hidden_size\n",
    "        # outputs就是所有时间步的结果，所以大小为(BATCH_SIZE, seq_len, hidden_size)\n",
    "        # 如果有多层，那么outputs中是最后一层的结果\n",
    "        # hidden为最后一个时间步的结果，如果有多层，那么hidden为每层最后一个时间步的结果\n",
    "        outputs, hidden = self.rnn(embedded, hidden)\n",
    "        # 本来outputs是(1, 10, 128)，out就变成了(1, 10, 109)  109是词汇表的大小，表示10个词分别预测出了10个词，预测的是每个词的概率\n",
    "        out = self.out_linear(outputs)\n",
    "\n",
    "        return out, hidden\n",
    "\n",
    "\n",
    "# 初始化模型\n",
    "model = CharGRU(vocab_size, HIDDEN_SIZE)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2348428ce74982e4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-23T02:16:22.779698Z",
     "start_time": "2025-05-23T02:15:59.863122Z"
    }
   },
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 36\u001b[0m\n\u001b[1;32m     32\u001b[0m             avg_loss \u001b[38;5;241m=\u001b[39m total_loss \u001b[38;5;241m/\u001b[39m \u001b[38;5;28mlen\u001b[39m(dataloader)\n\u001b[1;32m     33\u001b[0m             \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mEpoch [\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepoch\u001b[38;5;250m \u001b[39m\u001b[38;5;241m+\u001b[39m\u001b[38;5;250m \u001b[39m\u001b[38;5;241m1\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m/\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepochs\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m], Loss: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mavg_loss\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.4f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m---> 36\u001b[0m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataloader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mNUM_EPOCHS\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[13], line 8\u001b[0m, in \u001b[0;36mtrain\u001b[0;34m(model, dataloader, epochs)\u001b[0m\n\u001b[1;32m      4\u001b[0m total_loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m inputs, targets \u001b[38;5;129;01min\u001b[39;00m dataloader:\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;66;03m# 前向传播\u001b[39;00m\n\u001b[0;32m----> 8\u001b[0m     outputs, _ \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     10\u001b[0m     \u001b[38;5;66;03m# 计算损失\u001b[39;00m\n\u001b[1;32m     11\u001b[0m     \u001b[38;5;66;03m# 用每个时间步的输出和每个时间步的标签进行比较，并计算损失\u001b[39;00m\n\u001b[1;32m     12\u001b[0m     \u001b[38;5;66;03m# outputs的大小是(1, 10, 109)， targets的大小是(1, 10)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;66;03m# outputs.view(-1, vocab_size)表示将outputs变成(10, 109)，将三维变两维，一整个batch一起计算损失\u001b[39;00m\n\u001b[1;32m     14\u001b[0m     loss \u001b[38;5;241m=\u001b[39m criterion(\n\u001b[1;32m     15\u001b[0m         outputs\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, vocab_size),  \u001b[38;5;66;03m# (batch_size*seq_length, vocab_size)\u001b[39;00m\n\u001b[1;32m     16\u001b[0m         targets\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)  \u001b[38;5;66;03m# (batch_size*seq_length)\u001b[39;00m\n\u001b[1;32m     17\u001b[0m     )\n",
      "File \u001b[0;32m~/miniconda3/envs/mini-gpt/lib/python3.8/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/mini-gpt/lib/python3.8/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[11], line 35\u001b[0m, in \u001b[0;36mCharGRU.forward\u001b[0;34m(self, x, hidden)\u001b[0m\n\u001b[1;32m     33\u001b[0m outputs, hidden \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrnn(embedded, hidden)\n\u001b[1;32m     34\u001b[0m \u001b[38;5;66;03m# 本来outputs是(1, 10, 128)，out就变成了(1, 10, 109)  109是词汇表的大小，表示10个词分别预测出了10个词，预测的是每个词的概率\u001b[39;00m\n\u001b[0;32m---> 35\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241m.\u001b[39mout_linear(outputs)\n\u001b[1;32m     37\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out, hidden\n",
      "Cell \u001b[0;32mIn[11], line 35\u001b[0m, in \u001b[0;36mCharGRU.forward\u001b[0;34m(self, x, hidden)\u001b[0m\n\u001b[1;32m     33\u001b[0m outputs, hidden \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrnn(embedded, hidden)\n\u001b[1;32m     34\u001b[0m \u001b[38;5;66;03m# 本来outputs是(1, 10, 128)，out就变成了(1, 10, 109)  109是词汇表的大小，表示10个词分别预测出了10个词，预测的是每个词的概率\u001b[39;00m\n\u001b[0;32m---> 35\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241m.\u001b[39mout_linear(outputs)\n\u001b[1;32m     37\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out, hidden\n",
      "File \u001b[0;32m_pydevd_bundle/pydevd_cython_darwin_38_64.pyx:1187\u001b[0m, in \u001b[0;36m_pydevd_bundle.pydevd_cython_darwin_38_64.SafeCallWrapper.__call__\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m_pydevd_bundle/pydevd_cython_darwin_38_64.pyx:627\u001b[0m, in \u001b[0;36m_pydevd_bundle.pydevd_cython_darwin_38_64.PyDBFrame.trace_dispatch\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m_pydevd_bundle/pydevd_cython_darwin_38_64.pyx:1103\u001b[0m, in \u001b[0;36m_pydevd_bundle.pydevd_cython_darwin_38_64.PyDBFrame.trace_dispatch\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m_pydevd_bundle/pydevd_cython_darwin_38_64.pyx:1061\u001b[0m, in \u001b[0;36m_pydevd_bundle.pydevd_cython_darwin_38_64.PyDBFrame.trace_dispatch\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/Applications/PyCharm.app/Contents/plugins/python-ce/helpers/jupyter_debug/pydev_jupyter_plugin.py:171\u001b[0m, in \u001b[0;36mstop\u001b[0;34m(plugin, pydb, frame, event, args, stop_info, arg, step_cmd)\u001b[0m\n\u001b[1;32m    169\u001b[0m     frame \u001b[38;5;241m=\u001b[39m suspend_jupyter(main_debugger, thread, frame, step_cmd)\n\u001b[1;32m    170\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m frame:\n\u001b[0;32m--> 171\u001b[0m         \u001b[43mmain_debugger\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdo_wait_suspend\u001b[49m\u001b[43m(\u001b[49m\u001b[43mthread\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mframe\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mevent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43marg\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    172\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m    173\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
      "File \u001b[0;32m/Applications/PyCharm.app/Contents/plugins/python-ce/helpers/pydev/pydevd.py:1220\u001b[0m, in \u001b[0;36mPyDB.do_wait_suspend\u001b[0;34m(self, thread, frame, event, arg, send_suspend_message, is_unhandled_exception)\u001b[0m\n\u001b[1;32m   1217\u001b[0m         from_this_thread\u001b[38;5;241m.\u001b[39mappend(frame_id)\n\u001b[1;32m   1219\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_threads_suspended_single_notification\u001b[38;5;241m.\u001b[39mnotify_thread_suspended(thread_id, stop_reason):\n\u001b[0;32m-> 1220\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_do_wait_suspend\u001b[49m\u001b[43m(\u001b[49m\u001b[43mthread\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mframe\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mevent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43marg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msuspend_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfrom_this_thread\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/Applications/PyCharm.app/Contents/plugins/python-ce/helpers/pydev/pydevd.py:1235\u001b[0m, in \u001b[0;36mPyDB._do_wait_suspend\u001b[0;34m(self, thread, frame, event, arg, suspend_type, from_this_thread)\u001b[0m\n\u001b[1;32m   1232\u001b[0m             \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_mpl_hook()\n\u001b[1;32m   1234\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_internal_commands()\n\u001b[0;32m-> 1235\u001b[0m         \u001b[43mtime\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msleep\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m0.01\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1237\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcancel_async_evaluation(get_current_thread_id(thread), \u001b[38;5;28mstr\u001b[39m(\u001b[38;5;28mid\u001b[39m(frame)))\n\u001b[1;32m   1239\u001b[0m \u001b[38;5;66;03m# process any stepping instructions\u001b[39;00m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "def train(model, dataloader, epochs):\n",
    "    model.train()\n",
    "    for epoch in range(epochs):\n",
    "        total_loss = 0\n",
    "\n",
    "        for inputs, targets in dataloader:\n",
    "            # 前向传播\n",
    "            outputs, _ = model(inputs)\n",
    "\n",
    "            # 计算损失\n",
    "            # 用每个时间步的输出和每个时间步的标签进行比较，并计算损失\n",
    "            # outputs的大小是(1, 10, 109)， targets的大小是(1, 10)\n",
    "            # outputs.view(-1, vocab_size)表示将outputs变成(10, 109)，将三维变两维，一整个batch一起计算损失\n",
    "            loss = criterion(\n",
    "                outputs.view(-1, vocab_size),  # (batch_size*seq_length, vocab_size)\n",
    "                targets.view(-1)  # (batch_size*seq_length)\n",
    "            )\n",
    "\n",
    "            # 反向传播\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "\n",
    "            # 梯度裁剪防止爆炸\n",
    "            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "\n",
    "            optimizer.step()\n",
    "\n",
    "            total_loss += loss.item()\n",
    "\n",
    "        # 每20轮打印进度\n",
    "        if (epoch + 1) % 20 == 0:\n",
    "            avg_loss = total_loss / len(dataloader)\n",
    "            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}')\n",
    "\n",
    "\n",
    "train(model, dataloader, NUM_EPOCHS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a4e36b033918def",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_text(model, start_str, num_chars, temperature=0.8):\n",
    "    model.eval()\n",
    "    chars = [ch for ch in start_str]\n",
    "    input_seq = torch.LongTensor([word_to_idx[ch] for ch in chars[-SEQ_LENGTH:]])\n",
    "    hidden = None\n",
    "\n",
    "    for _ in range(num_chars):\n",
    "        # 输入形状调整\n",
    "        batch_input = input_seq.unsqueeze(0)  # (1, seq_len)\n",
    "\n",
    "        # 前向传播\n",
    "        with torch.no_grad():\n",
    "\n",
    "            # output中包含了每个时间步的输出，推理预测时，只需要取最后一个时间步的输出即可，比如输入“鹰击”，相当于有两个时间步，但是我们只需要第2个时间步的输出，而输出是词汇表中各个词的概率\n",
    "            # 而hidden表示隐藏层，在推理预测时，因为我们会连续预测，外层有一个for循环，所以hidden需要保存，以便下一次循环使用\n",
    "            output, hidden = model(batch_input, hidden)\n",
    "            last_output = output[0, -1, :]  # 最后时间步的输出\n",
    "\n",
    "        # 应用温度采样\n",
    "        # last_output / temperature，相当于将last_output缩小，比如[8,2,2] / 2 = [4,1,1]，使得三个选项对应的数字之间的差别变小了\n",
    "        # 当然如果temperature<1，那么就是放大差别，比如[8,2,2] / 0.5 = [16,4,4]\n",
    "        # probs为做了softmax之后的概率\n",
    "        probs = torch.softmax(last_output / temperature, dim=-1)\n",
    "\n",
    "        # 多项式采样，probs是一个概率，比如是[0.3,0.2,0.5]，那么就是从0,1,2中随机选一个，那么2被选中的概率就是50%\n",
    "        # 谁的概率大，随被采样的概率就大\n",
    "        char_idx = torch.multinomial(probs, 1).item()\n",
    "\n",
    "        # 更新输入序列\n",
    "        chars.append(idx_to_word[char_idx])\n",
    "        input_seq = torch.cat((input_seq[1:], torch.LongTensor([char_idx])))\n",
    "\n",
    "    return ''.join(chars)\n",
    "\n",
    "\n",
    "# 20表示预测20次, temperature越大，越随机\n",
    "print(generate_text(model, \"鹰击\", 20, temperature=0.1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58e0e8ea40651d38",
   "metadata": {},
   "source": [
    "效果确实比RNN好"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
